import numpy as np
from matplotlib import pyplot as plt

class FrenquencyRank:
    def __init__(self, img, gt, rank_len=10, type=True):
        self.rank_len = rank_len
        self.rank = []
        self.img = img / 255
        self.gt = gt
        self.type = type

    def add(self, adv, skl, loss, params):
        self.rank.insert(0, {
            'loss': round(loss, 3),
            'adv': adv / 255,
            'skl': skl,
            'params': params,
        })
        if len(self.rank) > self.rank_len:
            self.rank.pop()

    def check(self, adv, skl, loss, params):
        if len(self.rank) == 0:
            flag = True
        elif self.rank[0]['loss'] <= loss and self.type:
            flag = True
        else:
            flag = False

        if flag:
            self.add(adv, skl, loss, params)

    def best(self):
        return self.rank[0]

    def draw(self, save_dir, show=False):
        rank = self.rank
        length = len(rank)

        plt.figure()
        plt.subplot(2, length+1, 1)
        plt.imshow(self.img)
        plt.title('origin')
        plt.xticks([])
        plt.yticks([])
        plt.axis('off')

        if self.gt is not None:
            plt.subplot(2, length+1, length+2)
            plt.imshow(self.gt)
            plt.title('GT')
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')

        for i in range(length):
            item = rank[i]

            loss = item['loss']
            adv = item['adv']
            skl = item['skl']
            params = item['params']

            plt.subplot(2, length+1, i+2)
            plt.imshow(adv)
            plt.title('loss:' + str(loss))
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')

            if skl is not None:
                plt.subplot(2, length+1, length+i+3)
                plt.imshow(skl)
                plt.title('loss:' + str(loss))
                plt.xticks([])
                plt.yticks([])
                plt.axis('off')

        if show:
            plt.show()

        # plt.savefig(save_dir, dpi=300, bbox_inches='tight')
        